package org.jlambda.util;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicMarkableReference;
import java.util.concurrent.atomic.AtomicReference;


public class TernaryTree<V> implements Map<String, V>{
	static enum MatchedOn {equal, less, more};
	
	public interface Matches {
		boolean matches(String search, int pos, String currentNodePath);
	}
	
	public static final Matches equals = new TernaryTree.Matches(){
		public boolean matches(String search, int pos, String currentNodePath) {
			return search.equals(currentNodePath);
		}
	};
	
	/**
	 * Maintains a stack of visited nodes.  Pop's don't actually remove items from the stack, but merely
	 * reduce the stackPos.  This information is used to see where the tree was last visited. 
	 * 
	 * 
	 * @author Chris
	 *
	 */
	protected class Search {
		private String search = null;
		private StringBuilder path = new StringBuilder();
		private int pos = 0;
		private int stackPos = 0;
		
		// is it going up the tree?  make sure we go equals or more !!!
		private boolean up = false;
		private Matches matches = null;
		
		protected boolean isSearching(){
			return search != null;
		}
		
		public void goingUp(){
			up = true;
		}
		public void goingDown(){
			up = false;
		}
		
		public boolean isGoingUp(){
			return up;
		}
		
		protected boolean isLength(){
			if (search == null ) {
				return false;
			}
			return pos == search.length();
		}
		
		protected Search(){}
		protected Search(String searchString, Matches matches){
			search = searchString;		
			this.matches = matches;
		}
		
		/**
		 * Does the equal == null, if searching check the splichar is good, then try the matches.
		 * 
		 * If the next node in the stack (the one we have just come from if we are going up the tree) 
		 * is the equals then return false;
		 * 
		 * @param res
		 */
		protected boolean matchesEqual(Node res){
			if (res.equal.get() == null){
				return false;
			}
			
			if (stackPos + 1 < stack.size()){
				if (stack.get(stackPos).node == res.equal.get()) {
					return false;
				}
			}
			
			if (matches == null) {
				return true;
			}
			
			if (search == null){
				return true;
			}
			if (pos >-1 && search.length()> pos && search.charAt(pos) == res.splitChar) {
				return true;
			}
			
			return matches.matches(search, pos, path.toString() + res.splitChar);
		}
		
		
		protected boolean matchesValue(Node res){
			if (matches == null) {
				return true;
			}
			
			if (search == null){
				return true;
			}
			
			return matches.matches(search, pos, path.toString() + res.splitChar);
		}
		
		protected boolean matchesMore(Node res){
			if (res.more.get() == null){
				return false;
			}

			if (stackPos + 1 < stack.size()){
				if (stack.get(stackPos).node == res.more.get()) {
					return false;
				}
			}
			
			if (matches == null) {
				return true;
			}
			
			if (search == null){
				return true;
			}
			if (pos>-1 && search.length()> pos && res.splitChar < search.charAt(pos) ) {
				return true;
			}
			
			return matches.matches(search, pos, path.toString() + res.more.get().splitChar);
		}
		protected boolean matchesLess(Node res){
			if (res.less.get() == null){
				return false;
			}
			
			if (stackPos + 1 < stack.size()){
				if (stack.get(stackPos).node == res.less.get()) {
					return false;
				}
			}
			
			if (matches == null) {
				return true;
			}
			
			if (search == null){
				return true;
			}
			// The pos is now 1 but the splitChar is 0 f vs fa, so f > a and its true
			if (pos>-1 && search.length()> pos && res.splitChar > search.charAt(pos) ) {
				return true;
			}
			
			return matches.matches(search, pos, path.toString() + res.less.get().splitChar);
		}
		
		protected class pair {
			public Node node = null;
			public MatchedOn on = null;
			public pair(Node pos, MatchedOn mon){
				node = pos;
				on = mon;
			}
			@Override
			public String toString() {
				return node.toString() + " " + on.toString();
			}
		}
		
		ArrayList<pair> stack = new ArrayList<pair>();
		
		/**
		 * called when using equals
		 */
		protected void push(Node node, MatchedOn on){
			if (on == MatchedOn.equal) {
				pos++;
				path.append(node.splitChar);
			}
			
			if (i_isLast(node)) {
				// overwrite it
				stack.remove( --stackPos );
			}
			stack.add(stackPos ++, new pair(node, on));
		}
		
		/**
		 * going up tree / not equals down
		 */ 
		protected pair pop(){
			if (stackPos == 0){
				return null;
			}
			
			pair r = stack.get( --stackPos );
			if (r.on == MatchedOn.equal) {
				path.deleteCharAt(--pos);
			}
			last = r;
			return r;
		}
		
		protected String getCurrentPath(){
			return path.toString();
		}
		
		@Override
		public String toString() {
			return path.toString() + ":" + ((search == null)?"":search) +":"+pos;
		}

		// is it really last, more for pushing
		private boolean i_isLast(Node current) {
			if (stack.size() == 0) {
				return false;
			}
			if (stackPos - 1 < stack.size() && 0 < stackPos){
				if (stack.get(stackPos-1).node == current) {
					return true;
				}
			}
			return false;
		}
			
		public boolean isLast(Node current) {
			if (last != null && last.node == current){
				return true;
			}
			return i_isLast(current);
		}

		pair last = null;
		public void setLast(Node aboutToReturn, MatchedOn on){
			last = new pair(aboutToReturn, on);
		}
		
		public pair last() {
			return last;
		}

	}
	
	AtomicInteger size = new AtomicInteger(0);
	
	protected class Node {
		final protected char splitChar;
		protected Node(char val){
			splitChar = val;
		}
		final AtomicReference<Node> less = new AtomicReference<Node>();
		final AtomicReference<Node> equal = new AtomicReference<Node>();
		final AtomicReference<Node> more = new AtomicReference<Node>();
		final AtomicMarkableReference<V> value = new AtomicMarkableReference<V>(null, false);
		
		protected V getValue(){
			return value.getReference();
		}
		
		@Override
		public int hashCode() {
			int code = 37;
			code = code * 31 + splitChar;
			if (less.get() != null){
				code = code * 31 + less.get().splitChar;
			}
			if (more.get() != null){
				code = code * 31 + more.get().splitChar;
			}
			if (equal.get() != null){
				code = code * 31 + equal.get().splitChar;
			}

			return code;
		}
		
		@Override
		public boolean equals(Object obj) {
			if (obj instanceof TernaryTree.Node) {
				Node lhs = (Node) obj;
				if (lhs.splitChar == splitChar) {
					if (lhs.less.get() == less.get()){
						if (lhs.equal.get() == equal.get()) {
							if (lhs.more.get() == more.get()) {
									return true;
							}
						}
					}
				}
			}
			return false;
		}
		
		@Override
		public String toString() {
			return splitChar +":"+( (value.getReference() == null)?"":value.getReference() );
		}
	}

	protected AtomicReference<Node> root = new AtomicReference<Node>();
	
	// returns null if the node could not be added
	// returns the new node if added, or the currently inserted node
	protected Node checkAddNode(AtomicReference<Node> ref, char nextPos){
		Node current = null;
		if (ref.get() == null ){
			// create new one
			Node eq = new Node(nextPos);
			if (ref.compareAndSet(null, eq)){
				current = ref.get();
			} else {
				// didn't set it in which case we need to retry
				return null;
			}
		} else {
			current = ref.get();
		}
		return current;
	}
	
	protected V insert(String key, V val, V expected, boolean ifabsent){
		Node current = root.get();
		int pos = 0;
		if (current == null) {
			current = new Node(key.charAt(pos));
			if (!root.compareAndSet(null, current)){
				current = root.get();
			}
		}
		while(true) {
			char cpos = key.charAt(pos);
			if (pos == key.length()-1 && current.splitChar == cpos ) {
				if (!ifabsent) {
					// set it anyway
					V res = current.value.getReference();
					boolean isThere = current.value.isMarked();
					current.value.set(val,true);
//					current.key = key;
					if (!isThere) {
						size.incrementAndGet();
					}
					return res;
				}
				if (current.value.compareAndSet(expected, val, false, true)){
					// wasn't there
					size.incrementAndGet();
					return (ifabsent)?null:val;
				} else {
					// was already there, as such we should return the existing one
					return current.value.getReference();
				}
			} else if (pos == key.length()) {
				throw new RuntimeException("no idea why it didn't insert");
			} else {
				if (current.splitChar == cpos){
					Node res = checkAddNode(current.equal, key.charAt(pos+1));
					if (res != null) {
						pos++;
						current = res;
					}
				} else if (current.splitChar < cpos){
					// if its empty then we should add it
					Node res = checkAddNode(current.more, cpos);
					if (res != null) {
						current = res;
					}
				} else if (current.splitChar > cpos){
					// if its empty then we should add it
					Node res = checkAddNode(current.less, cpos);
					if (res != null) {
						current = res;
					}
				}
			}
		}
	}
	
	protected Node getNode(Search search){
		Node current = null;
		MatchedOn on = MatchedOn.more;
		Search.pair last = search.last();
		if (last == null) {
			current = root.get();
		} else {
			current = last.node;
			on = last.on;
		}
		boolean up = search.isGoingUp();
		
		if (current == null) {
			return null;
		}
		while(current != null) {
			// evaluate before we get into ifs
			boolean islast = search.isLast(current);
			up = search.isGoingUp();
			
			// ifs are arranged to reduce calls on matchers
			
			if (!up && search.matchesLess(current)){
				on = MatchedOn.less;
				search.push(current, MatchedOn.less);
				current = current.less.get();
			} else if (!islast && (!up || (up && on == MatchedOn.less) ) && current.value.isMarked() && search.matchesValue(current) ){
				search.setLast(current, on);
				return current;
			} else if ((!up || (up && on == MatchedOn.less) ) && search.matchesEqual(current)){
				search.push(current, MatchedOn.equal);
				search.goingDown();
				on = MatchedOn.equal;
				current = current.equal.get();
			
			} else if ((!up || (up && on == MatchedOn.equal) ) && search.matchesMore(current)){
				on = MatchedOn.more;
				search.goingDown();
				search.push(current, MatchedOn.more);
				current = current.more.get();
			} else {
				// we've reached the end of this tree, time to pop
				search.goingUp();
				
				last = search.pop();
				if (last == null) {
					if (on == MatchedOn.more || current == root.get()) {
						return null; // were are at root after going root.more
					}
					current = root.get();
				} else {
					current = last.node;
					on = last.on;
				}
				
				// as we are on the up, we must check for the value at this node
				if (current.value.isMarked() && on == MatchedOn.less && search.matchesValue(current) ) {
					search.setLast(current, on);
					return current;
				}
			}
		}
		return null;
	}

	
	protected Node getNode(String key){
		Node current = root.get();
		int pos = 0;
		if (current == null) {
			return null;
		}
		while(current != null) {
			char cpos = key.charAt(pos);
			if (pos == key.length()-1 && current.splitChar == cpos) {
				// must be here
				if (!current.value.isMarked()) {
					// key is there, but its been deleted
					return null;
				} else {
					// return current value
					return current;
				}
			} else {
				if (current.splitChar == cpos){
					if (current.equal == null) {
						// not a valid key
						return null;
					}
					pos++;
					current = current.equal.get();
				} else if (current.splitChar < cpos){
					if (current.more == null) {
						// not a valid key
						return null;
					}
					current = current.more.get();
				} else if (current.splitChar > cpos){
					if (current.less == null) {
						// not a valid key
						return null;
					}
					current = current.less.get();
				}
			}
		}
		return null;
	}
	
	protected V get(String key){
		Node n = getNode(key); 
		return (n == null) ? null : n.value.getReference();
	}
	
	protected V remove(String key){
		Node node = getNode(key);
		if (node == null ){
			return null;
		}
		
		V res = node.value.getReference();
		node.value.set(null, false); // wipe it out
		size.decrementAndGet();
		return res;
	}
	
	class NodeIterator implements Iterator<Map.Entry<String, V> > {

		Node current = null;
		
		boolean eol = false;
		private Search search;
		
		NodeIterator(Search search) {
			this.search = search;
			current = getNode(search);
			if (current == null) {
				eol = true;
			}
		}
		
		public boolean hasNext() {
			if (eol) {
				return false;
			}
			
			if (current != null) {
				return true;
			}
			return false;
		}
		
		class entry implements Entry<String, V> {

			final Node node;
			final String key;
			entry(Node node, String key){
				this.node = node;
				this.key = key;
			}
			
			public String getKey() {
				return key;
			}

			public V getValue() {
				return node.value.getReference();
			}

			public V setValue(V value) {
				return null;
			}
			
			@Override
			public String toString() {
				return key+","+getValue();
			}
		}

		public Entry<String, V> next() {
			if (eol) {
				throw new NoSuchElementException("no more entries");
			}
			
			Entry<String, V> res = null;
			if (current != null){
				res = new entry(current, search.getCurrentPath() + current.splitChar);
				current = getNode(search);
			} else {
				eol = true;
				throw new NoSuchElementException("no more entries");
			}
			
			return res;
		}

		public void remove() {
			
		}
		
	}
	
	class entrySet implements Set<Entry<String,V>>{

		public int size() {
			return TernaryTree.this.size();
		}

		public boolean isEmpty() {
			return TernaryTree.this.isEmpty();
		}

		public boolean contains(Object o) {
			return false;
		}

		public Iterator<Entry<String, V>> iterator() {
			return new NodeIterator(new Search());
		}

		public Object[] toArray() {
			return null;
		}

		public <T> T[] toArray(T[] a) {
			return null;
		}

		public boolean add(Entry<String, V> o) {
			TernaryTree.this.put(o.getKey(), o.getValue());
			return true;
		}

		public boolean remove(Object o) {
			return false;
		}

		public boolean containsAll(Collection<?> c) {
			return false;
		}

		public boolean addAll(Collection<? extends Entry<String, V>> c) {
			for (Entry<? extends String, ? extends V> entry : c) {
				put(entry.getKey(), entry.getValue());
			}
			return true;
		}

		public boolean retainAll(Collection<?> c) {
			return false;
		}

		public boolean removeAll(Collection<?> c) {
			return false;
		}

		public void clear() {
			TernaryTree.this.clear();	
		}
		
	}
	
	public V putIfAbsent(String key, V value) {
		return insert(key, value, null, true);
	}

	public boolean remove(Object key, Object value) {
		Node node = getNode(key.toString());
		if (node == null) {
			return false;
		}
		
		// try wipe it out
		if (node.value.compareAndSet((V) value, null,true, false)){
			size.decrementAndGet();
			return true;
		}
		return false;
	}

	public boolean replace(String key, V oldValue, V newValue) {
		Node node = getNode(key.toString());
		if (node == null) {
			return false;
		}
		
		// try and replace it
		if (node.value.compareAndSet((V) oldValue, (V)newValue,true, true)){
			return true;
		}
		return false;
	}

	public V replace(String key, V value) {
		Node node = getNode(key.toString());
		if (node == null) {
			return null;
		}
		
		// get value
		V res = node.value.getReference();
		node.value.set(value,true);
		return res;
	}

	public int size() {
		return size.get();
	}

	public boolean isEmpty() {
		return root.get() == null;
	}

	public boolean containsKey(Object key) {
		return getNode(key.toString()) != null;
	}

	public boolean containsValue(Object value) {
		return false;
	}

	public V get(Object key) {
		return get(key.toString());
	}

	public V put(String key, V value) {
		return insert(key, value, null, false);
	}

	public V remove(Object key) {
		return remove(key.toString());
	}

	public void clear() {
		int c = size.get();
		root.set(null);
		size.addAndGet(-c); // note this must be done, as between the set and the addAndGet it could have nodes added
	}
	
	public Set<Entry<String,V>> entrySet(){
		return new entrySet();
	}

	public void putAll(Map<? extends String, ? extends V> t) {
		for (Entry<? extends String, ? extends V> entry : t.entrySet()) {
			put(entry.getKey(), entry.getValue());
		}
	}
	
	class keys implements Set<String>{

		public int size() {
			return TernaryTree.this.size();
		}

		public boolean isEmpty() {
			return TernaryTree.this.isEmpty();
		}

		public boolean contains(Object o) {
			return containsKey(o.toString());
		}

		class keyIterator implements Iterator<String>{

			NodeIterator itr = new NodeIterator(new Search());
			
			public boolean hasNext() {
				return itr.hasNext();
			}

			public String next() {
				return itr.next().getKey();
			}

			public void remove() {
			}
		}
		
		public Iterator<String> iterator() {
			return new keyIterator();
		}

		public Object[] toArray() {
			return null;
		}

		public <T> T[] toArray(T[] a) {
			return null;
		}

		public boolean add(String o) {
			return true;
		}

		public boolean remove(Object o) {
			return false;
		}

		public boolean containsAll(Collection<?> c) {
			return false;
		}

		public boolean addAll(Collection<? extends String> c) {
			return false;
		}

		public boolean retainAll(Collection<?> c) {
			return false;
		}

		public boolean removeAll(Collection<?> c) {
			return false;
		}

		public void clear() {
			TernaryTree.this.clear();
		}
		
	}

	public Set<String> keySet() {
		return new keys();
	}

	class values implements Collection<V>{

		public int size() {
			return TernaryTree.this.size();
		}

		public boolean isEmpty() {
			return TernaryTree.this.isEmpty();
		}

		public boolean contains(Object o) {
			return containsValue(o);
		}

		class valueIterator implements Iterator<V>{

			NodeIterator itr = new NodeIterator(new Search());
			
			public boolean hasNext() {
				return itr.hasNext();
			}

			public V next() {
				V val = itr.next().getValue();
				return val;
			}

			public void remove() {
			}
			
		}
		
		public Iterator<V> iterator() {
			return new valueIterator();
		}

		public Object[] toArray() {
			return null;
		}

		public <T> T[] toArray(T[] a) {
			return null;
		}

		public boolean add(V o) {
			return false;
		}

		public boolean remove(Object o) {
			return false;
		}

		public boolean containsAll(Collection<?> c) {
			return false;
		}

		public boolean addAll(Collection<? extends V> c) {
			return false;
		}

		public boolean removeAll(Collection<?> c) {
			return false;
		}

		public boolean retainAll(Collection<?> c) {
			return false;
		}

		public void clear() {
			TernaryTree.this.clear();
		}
		
	}
	
	public Collection<V> values() {
		return new values();
	}

	public Iterator<Map.Entry<String,V>> prefixedBy(final String prefix){
		Matches matches = new Matches(){
			public boolean matches(String search, int pos, String currentNodePath) {
				return (currentNodePath.startsWith(prefix));
			}
		};
		Search search = new Search(prefix, matches);
		return new NodeIterator(search);		
	}

	/**
	 * General form for passing in custom matching functions.
	 * 
	 * @param searchString
	 * @param matches
	 * @return
	 */
	public Iterator<Map.Entry<String,V>> iterator(final String searchString, final Matches matches){
		Search search = new Search(searchString, matches);
		return new NodeIterator(search);		
	}
	
	public Iterator<Map.Entry<String,V>> lessThanEqual(final String prefix){
		Matches matches = new Matches(){
			public boolean matches(String search, int pos, String currentNodePath) {
				return currentNodePath.compareTo(search) <= 0;
			}
		};
		Search search = new Search(prefix, matches);
		return new NodeIterator(search);		
	}
	public Iterator<Map.Entry<String,V>> greaterThanEqual(final String prefix){
		Matches matches = new Matches(){
			public boolean matches(String search, int pos, String currentNodePath) {
				return currentNodePath.compareTo(search) >= 0;
			}
		};
		Search search = new Search(prefix, matches);
		return new NodeIterator(search);		
	}
}
